import asyncio
import websockets
import json
import uuid
from typing import Dict, List

class GomokuServer:
    def __init__(self):
        self.games: Dict[str, dict] = {}  # 存储游戏信息
        self.players: Dict[str, dict] = {}  # 存储玩家信息

    async def handle_connection(self, websocket):
        """处理WebSocket连接"""
        player_id = str(uuid.uuid4())
        self.players[player_id] = {
            "websocket": websocket,
            "game_id": None
        }
        
        try:
            async for message in websocket:
                await self.process_message(player_id, message)
        except websockets.exceptions.ConnectionClosed:
            pass
        finally:
            # 玩家断开连接时的清理工作
            await self.handle_disconnect(player_id)

    async def process_message(self, player_id, message):
        """处理来自客户端的消息"""
        try:
            data = json.loads(message)
            action = data.get("action")
            
            if action == "create_game":
                await self.create_game(player_id, data)
            elif action == "join_game":
                await self.join_game(player_id, data)
            elif action == "make_move":
                await self.make_move(player_id, data)
            elif action == "restart_game":
                await self.restart_game(player_id)
        except json.JSONDecodeError:
            pass  # 忽略无效的JSON消息

    async def create_game(self, player_id, data):
        """创建新游戏"""
        game_id = str(uuid.uuid4())[:8]
        self.games[game_id] = {
            "id": game_id,
            "players": [player_id],  # 先添加创建者
            "board": [[0 for _ in range(15)] for _ in range(15)],
            "current_player": 1,  # 1代表黑子，2代表白子，3代表红子
            "status": "waiting",  # waiting, playing, finished
            "winner": None,
            "player_count": 3  # 支持3人游戏
        }
        
        # 更新玩家信息
        self.players[player_id]["game_id"] = game_id
        self.players[player_id]["player_number"] = 1  # 创建者是黑子(1号玩家)
        
        # 发送游戏创建成功的消息
        await self.players[player_id]["websocket"].send(json.dumps({
            "action": "game_created",
            "game_id": game_id,
            "player_number": 1
        }))

    async def join_game(self, player_id, data):
        """加入游戏"""
        game_id = data.get("game_id")
        
        # 检查游戏是否存在
        if game_id not in self.games:
            await self.players[player_id]["websocket"].send(json.dumps({
                "action": "error",
                "message": "游戏不存在"
            }))
            return
        
        game = self.games[game_id]
        
        # 检查游戏是否已满
        if len(game["players"]) >= game["player_count"]:
            await self.players[player_id]["websocket"].send(json.dumps({
                "action": "error",
                "message": "游戏已满"
            }))
            return
        
        # 分配玩家编号
        player_number = len(game["players"]) + 1
        # 加入游戏
        game["players"].append(player_id)
        
        # 更新玩家信息
        self.players[player_id]["game_id"] = game_id
        self.players[player_id]["player_number"] = player_number
        
        # 如果达到3人，开始游戏
        if len(game["players"]) == game["player_count"]:
            game["status"] = "playing"
            
        # 通知所有玩家游戏状态更新
        for i, p_id in enumerate(game["players"]):
            p_number = self.players[p_id]["player_number"]
            await self.players[p_id]["websocket"].send(json.dumps({
                "action": "game_updated",
                "game_id": game_id,
                "player_number": p_number,
                "players_joined": len(game["players"]),
                "total_players": game["player_count"],
                "player_list": [
                    {"id": game["players"][j], "number": self.players[game["players"][j]]["player_number"]} 
                    for j in range(len(game["players"]))
                ]
            }))
            
        # 如果游戏开始，通知所有玩家
        if game["status"] == "playing":
            for p_id in game["players"]:
                p_number = self.players[p_id]["player_number"]
                await self.players[p_id]["websocket"].send(json.dumps({
                    "action": "game_started",
                    "game_id": game_id,
                    "player_number": p_number,
                    "board": game["board"],
                    "current_player": game["current_player"]
                }))

    async def make_move(self, player_id, data):
        """处理玩家落子"""
        game_id = self.players[player_id].get("game_id")
        
        # 检查玩家是否在游戏内
        if not game_id or game_id not in self.games:
            return
        
        game = self.games[game_id]
        
        # 检查是否轮到该玩家
        player_number = self.players[player_id]["player_number"]
        if game["current_player"] != player_number:
            return
        
        row = data.get("row")
        col = data.get("col")
        
        # 检查位置是否有效
        if not (0 <= row < 15 and 0 <= col < 15):
            return
        
        # 检查位置是否为空
        if game["board"][row][col] != 0:
            return
        
        # 落子
        game["board"][row][col] = player_number
        
        # 检查是否获胜
        winner = self.check_winner(game["board"], row, col, player_number)
        
        # 切换当前玩家（除非游戏已结束）
        if not winner:
            game["current_player"] = (game["current_player"] % game["player_count"]) + 1  # 1->2, 2->3, 3->1
        
        # 通知所有玩家落子信息
        for p_id in game["players"]:
            await self.players[p_id]["websocket"].send(json.dumps({
                "action": "move_made",
                "row": row,
                "col": col,
                "player": player_number,
                "current_player": game["current_player"],
                "winner": winner,
                "board": game["board"]
            }))
            
            # 如果有获胜者，更新游戏状态
            if winner:
                game["status"] = "finished"
                game["winner"] = winner

    async def restart_game(self, player_id):
        """重新开始游戏"""
        game_id = self.players[player_id].get("game_id")
        
        # 检查玩家是否在游戏内
        if not game_id or game_id not in self.games:
            return
        
        game = self.games[game_id]
        
        # 重置游戏状态
        game["board"] = [[0 for _ in range(15)] for _ in range(15)]
        game["current_player"] = 1  # 黑子先手
        game["status"] = "playing"
        game["winner"] = None
        
        # 通知所有玩家游戏重新开始
        for p_id in game["players"]:
            await self.players[p_id]["websocket"].send(json.dumps({
                "action": "game_restarted",
                "board": game["board"],
                "current_player": game["current_player"]
            }))

    def check_winner(self, board, row, col, player):
        """检查是否有玩家获胜"""
        directions = [
            [(0, 1), (0, -1)],   # 水平
            [(1, 0), (-1, 0)],   # 垂直
            [(1, 1), (-1, -1)],  # 主对角线
            [(1, -1), (-1, 1)]   # 副对角线
        ]
        
        for direction_pair in directions:
            count = 1  # 包含当前落子点
            
            # 向一个方向检查
            dx, dy = direction_pair[0]
            r, c = row + dx, col + dy
            while 0 <= r < 15 and 0 <= c < 15 and board[r][c] == player:
                count += 1
                r, c = r + dx, c + dy
            
            # 向相反方向检查
            dx, dy = direction_pair[1]
            r, c = row + dx, col + dy
            while 0 <= r < 15 and 0 <= c < 15 and board[r][c] == player:
                count += 1
                r, c = r + dx, c + dy
            
            # 如果连成五子，则当前玩家获胜
            if count >= 5:
                return player
        
        return None

    async def handle_disconnect(self, player_id):
        """处理玩家断开连接"""
        if player_id in self.players:
            game_id = self.players[player_id].get("game_id")
            
            # 从玩家列表中移除
            del self.players[player_id]
            
            # 如果玩家在游戏中，通知对手
            if game_id and game_id in self.games:
                game = self.games[game_id]
                # 通知游戏中的其他玩家
                for p_id in game["players"]:
                    if p_id != player_id and p_id in self.players:
                        try:
                            await self.players[p_id]["websocket"].send(json.dumps({
                                "action": "opponent_disconnected"
                            }))
                        except:
                            pass
                
                # 从游戏的玩家列表中移除
                if player_id in game["players"]:
                    game["players"].remove(player_id)
                
                # 如果游戏没有玩家了，删除游戏
                if len(game["players"]) == 0:
                    del self.games[game_id]

# 启动服务器
async def main():
    server = GomokuServer()
    
    print("五子棋服务器启动中...")
    print("服务器地址: ws://0.0.0.0:8765")
    
    async with websockets.serve(server.handle_connection, "0.0.0.0", 8765):
        await asyncio.Future()  # 运行直到被中断

if __name__ == "__main__":
    asyncio.run(main())